use ndarray::prelude::*;
use ndarray::{Array1, Array2};
use ndarray_linalg::Solve;
use std::collections::VecDeque;
use std::f64::consts::PI;
use rand::Rng;
use chrono::{NaiveDate, Duration};

// 数据结构：存储报废事件
#[derive(Debug, Clone)]
struct ScrapEvent {
    timestamp: f64, // 时间戳（从参考日期开始的天数）
    scrap_amount: f64, // 该时间点的报废数量
}

// 趋势+季节性模型
struct TrendSeasonalModel {
    trend_coefficients: Array1<f64>,    // [截距, 斜率]
    seasonal_coefficients: Array1<f64>, // 傅里叶系数
    period: f64,                        // 周期长度（天）
    fourier_terms: usize,               // 傅里叶项数
    last_train_time: f64,               // 最后训练时间
}

// 预测系统配置
struct ForecastConfig {
    max_history: usize,    // 最大历史数据点数
    min_history: usize,    // 最小历史数据点数
    retrain_interval: f64, // 重新训练间隔（天）
    data_decay_rate: f64,  // 数据衰减率
    period: f64,           // 季节性周期（天）
    fourier_terms: usize,  // 傅里叶项数
    min_interval: f64,     // 最小时间间隔阈值（天）
}

// 预测系统状态
struct ForecastingSystem {
    config: ForecastConfig,
    events: VecDeque<ScrapEvent>,       // 报废事件记录（按时间排序）
    cumulative_scrap: VecDeque<DataPoint>, // 累计报废量序列（用于训练）
    model: Option<TrendSeasonalModel>,  // 当前模型
    last_retrain_time: f64,             // 最后训练时间
    reference_date: NaiveDate,          // 参考日期（时间戳0点）
}

// 辅助结构：累计报废量数据点
#[derive(Debug, Clone)]
struct DataPoint {
    timestamp: f64,
    cumulative_scrap: f64,
}

impl TrendSeasonalModel {
    fn new() -> Self {
        Self {
            trend_coefficients: Array1::zeros(2),
            seasonal_coefficients: Array1::zeros(0),
            period: 365.0, // 默认年周期
            fourier_terms: 2,
            last_train_time: 0.0,
        }
    }

    // 使用加权最小二乘法拟合模型
    fn fit_weighted(
        times: &Array1<f64>,
        scrap_rates: &Array1<f64>,
        weights: &Array1<f64>,
        period: f64,
        fourier_terms: usize,
    ) -> Option<Self> {
        let n = times.len();
        let total_terms = 2 + 2 * fourier_terms; // 趋势(2) + 傅里叶(2*terms)
        
        if n < total_terms {
            return None; // 数据不足
        }
        
        let mut features = Array2::zeros((n, total_terms));
        let mut weighted_values = Array1::zeros(n);
        
        for i in 0..n {
            let t = times[i];
            let w = weights[i].sqrt();
            features[[i, 0]] = w * 1.0; // 常数项
            features[[i, 1]] = w * t;    // 时间项
            
            for k in 1..=fourier_terms {
                let angle = 2.0 * PI * k as f64 * t / period;
                features[[i, 2 + 2*k - 2]] = w * angle.sin(); // sin项
                features[[i, 2 + 2*k - 1]] = w * angle.cos(); // cos项
            }
            
            weighted_values[i] = w * scrap_rates[i];
        }
        
        let regularization: Array2<f64> = Array2::eye(total_terms) * 1e-6;
        let xt = features.t();
        let xtx = xt.dot(&features) + &regularization;
        let xty = xt.dot(&weighted_values);
        
        match xtx.solve(&xty) {
            Ok(all_coefficients) => {
                let trend_coefficients = all_coefficients.slice(s![0..2]).to_owned();
                let seasonal_coefficients = all_coefficients.slice(s![2..]).to_owned();
                Some(Self {
                    trend_coefficients,
                    seasonal_coefficients,
                    period,
                    fourier_terms,
                    last_train_time: *times.iter().max_by(|a, b| a.partial_cmp(b).unwrap()).unwrap_or(&0.0),
                })
            },
            Err(e) => {
                eprintln!("模型拟合失败: {:?}", e);
                None
            }
        }
    }

    // 预测指定时间点的报废率
    fn predict(&self, time: f64) -> f64 {
        let trend = self.trend_coefficients[0] + self.trend_coefficients[1] * time;
        let mut seasonal = 0.0;
        
        for k in 0..self.fourier_terms {
            let angle = 2.0 * PI * (k + 1) as f64 * time / self.period;
            seasonal += self.seasonal_coefficients[2*k] * angle.sin()
                     + self.seasonal_coefficients[2*k + 1] * angle.cos();
        }
        
        (trend + seasonal).max(0.0) // 确保非负
    }
}

impl ForecastConfig {
    fn default() -> Self {
        Self {
            max_history: 36,
            min_history: 3,
            retrain_interval: 30.0, // 每月重新训练
            data_decay_rate: 0.005, // 每日衰减率
            period: 365.0,          // 年周期
            fourier_terms: 2,
            min_interval: 0.1,      // 最小间隔0.1天
        }
    }
}

impl ForecastingSystem {
    fn new(reference_date: NaiveDate) -> Self {
        let config = ForecastConfig::default();
        let max_history = config.max_history;
        Self {
            config,
            events: VecDeque::with_capacity(max_history),
            cumulative_scrap: VecDeque::with_capacity(max_history),
            model: None,
            last_retrain_time: 0.0,
            reference_date,
        }
    }

    // 添加新报废事件（使用日期）
    fn add_scrap_event(&mut self, date: NaiveDate, scrap_amount: f64) {
        // 计算时间戳（从参考日期开始的天数）
        let days_since_ref = (date - self.reference_date).num_days() as f64;
        
        // 确保时间顺序
        if let Some(last_event) = self.events.back() {
            if days_since_ref < last_event.timestamp {
                eprintln!("错误：新事件日期 {} 早于最后事件日期", date);
                return;
            }
        }
        
        // 添加事件记录
        if self.events.len() >= self.config.max_history {
            self.events.pop_front();
        }
        self.events.push_back(ScrapEvent {
            timestamp: days_since_ref,
            scrap_amount,
        });
        
        // 重建累计报废序列
        self.rebuild_cumulative_series();
        
        // 检查是否需要重新训练
        if self.should_retrain(days_since_ref) {
            self.retrain(days_since_ref);
        }
    }
    
    // 重建累计报废序列
    fn rebuild_cumulative_series(&mut self) {
        self.cumulative_scrap.clear();
        let mut cumulative = 0.0;
        
        // 添加初始点（参考日期，报废量0）
        self.cumulative_scrap.push_back(DataPoint {
            timestamp: 0.0,
            cumulative_scrap: 0.0,
        });
        
        // 按时间排序事件
        let mut sorted_events: Vec<_> = self.events.iter().collect();
        sorted_events.sort_by(|a, b| a.timestamp.partial_cmp(&b.timestamp).unwrap());
        
        for event in sorted_events {
            cumulative += event.scrap_amount;
            self.cumulative_scrap.push_back(DataPoint {
                timestamp: event.timestamp,
                cumulative_scrap: cumulative,
            });
        }
    }

    // 判断是否需要重新训练
    fn should_retrain(&self, current_time: f64) -> bool {
        if self.cumulative_scrap.len() < self.config.min_history + 1 {
            return false;
        }
        current_time - self.last_retrain_time >= self.config.retrain_interval
    }

    // 重新训练模型
    fn retrain(&mut self, current_time: f64) {
        if self.cumulative_scrap.len() < 2 {
            return;
        }
        
        let n = self.cumulative_scrap.len() - 1;
        let mut times = Vec::with_capacity(n);
        let mut scrap_rates = Vec::with_capacity(n);
        let mut weights = Vec::with_capacity(n);
        
        for i in 1..self.cumulative_scrap.len() {
            let prev = &self.cumulative_scrap[i-1];
            let curr = &self.cumulative_scrap[i];
            
            let interval = curr.timestamp - prev.timestamp;
            if interval < self.config.min_interval {
                continue;
            }
            
            let scrap_amount = curr.cumulative_scrap - prev.cumulative_scrap;
            if scrap_amount <= 0.0 {
                continue;
            }
            
            let scrap_rate = scrap_amount / interval;
            let midpoint_time = (prev.timestamp + curr.timestamp) / 2.0;
            
            times.push(midpoint_time);
            scrap_rates.push(scrap_rate);
            
            // 计算权重（基于时间衰减和间隔长度）
            let time_weight = (-self.config.data_decay_rate * (current_time - curr.timestamp)).exp();
            let interval_weight = interval.sqrt();
            weights.push(time_weight * interval_weight);
        }
        
        if times.len() < self.config.min_history {
            return;
        }
        
        let times_arr = Array1::from(times.clone());
        let scrap_rates_arr = Array1::from(scrap_rates);
        let weights_arr = Array1::from(weights);
        
        if let Some(new_model) = TrendSeasonalModel::fit_weighted(
            &times_arr,
            &scrap_rates_arr,
            &weights_arr,
            self.config.period,
            self.config.fourier_terms,
        ) {
            self.model = Some(new_model);
            self.last_retrain_time = current_time;
            println!("模型训练成功，使用 {} 个数据点", times.len());
        } else {
            println!("模型训练失败");
        }
    }

    // 预测未来时间点的报废率
    fn predict_scrap_rate(&self, future_time: f64) -> Option<f64> {
        self.model.as_ref().map(|model| model.predict(future_time))
    }

    // 预测库存归零时间（天）
    fn predict_inventory_depletion(&self, current_inventory: f64, current_date: NaiveDate) -> Option<f64> {
        if current_inventory <= 0.0 {
            return Some(0.0); // 库存已耗尽
        }
        
        // 计算当前时间戳
        let current_time = (current_date - self.reference_date).num_days() as f64;
        
        // 获取模型预测或使用历史平均
        let default_scrap_rate = self.calculate_historical_average();
        let scrap_rate = self.predict_scrap_rate(current_time).unwrap_or(default_scrap_rate);
        
        // 如果没有有效报废率，返回None
        if scrap_rate <= 0.0 {
            return None;
        }
        
        // 简单估算：剩余库存 / 当前报废率
        let mut time_to_depletion = current_inventory / scrap_rate;
        
        // 使用模型进行精确积分预测（考虑报废率随时间变化）
        if let Some(model) = &self.model {
            let mut remaining = current_inventory;
            let mut time = current_time;
            let step = 0.1; // 时间步长（天）
            let max_time = current_time + 2.0 * time_to_depletion; // 最大预测时间
            
            while remaining > 0.0 && time < max_time {
                let rate = model.predict(time);
                if rate <= 0.0 {
                    break;
                }
                
                let consumption = rate * step;
                if consumption >= remaining {
                    let fraction = remaining / consumption;
                    time_to_depletion = (time + fraction * step) - current_time;
                    break;
                }
                
                remaining -= consumption;
                time += step;
            }
        }
        
        Some(time_to_depletion)
    }
    
    // 计算历史平均报废率
    fn calculate_historical_average(&self) -> f64 {
        if self.cumulative_scrap.len() < 2 {
            return 0.0;
        }
        
        let total_scrap = self.cumulative_scrap.back().unwrap().cumulative_scrap;
        let total_time = self.cumulative_scrap.back().unwrap().timestamp;
        
        if total_time > 0.0 {
            total_scrap / total_time
        } else {
            0.0
        }
    }
    
    // 获取总报废量
    fn total_scrap(&self) -> f64 {
        self.cumulative_scrap.back().map_or(0.0, |p| p.cumulative_scrap)
    }
}

// 示例使用
fn main() {
    // 设置参考日期（最早事件的日期）
    let reference_date = NaiveDate::from_ymd_opt(2024, 1, 1).unwrap();
    let mut system = ForecastingSystem::new(reference_date);
    
    // 添加报废事件（日期，报废数量）
    system.add_scrap_event(NaiveDate::from_ymd_opt(2024, 1, 1).unwrap(), 32.0);
    system.add_scrap_event(NaiveDate::from_ymd_opt(2024, 2, 2).unwrap(), 60.0);
    system.add_scrap_event(NaiveDate::from_ymd_opt(2024, 4, 5).unwrap(), 100.0);
system.add_scrap_event(NaiveDate::from_ymd_opt(2024, 5, 1).unwrap(), 50.0);
    
    println!("已添加 {} 个报废事件，总报废量: {}", 
        system.events.len(), system.total_scrap());
    
    // 用户输入当前库存和日期
    let current_inventory = 500.0;
    let current_date = NaiveDate::from_ymd_opt(2025, 10, 24).unwrap();
    
    // 预测库存归零时间（天）
    if let Some(days_to_depletion) = system.predict_inventory_depletion(current_inventory, current_date) {
        let depletion_date = current_date + Duration::days(days_to_depletion as i64);
        println!("\n当前库存 {} 件将在 {:.1} 天后耗尽", current_inventory, days_to_depletion);
        println!("预计耗尽日期: {}", depletion_date);
        
        // 预测未来报废率
        println!("\n未来报废率预测:");
        for days in [30, 60, 90] {
            let future_date = current_date + Duration::days(days);
            let future_time = (future_date - reference_date).num_days() as f64;
            if let Some(rate) = system.predict_scrap_rate(future_time) {
                println!("  {}天后 ({}) 预测报废率: {:.1}件/天", days, future_date, rate);
            }
        }
    } else {
        println!("\n无法预测库存耗尽时间");
    }
}